"""
根据实验记录撰写论文
"""

import re
import os
from translate import Translator

from .pubmed_searcher import PubMedSearcher
from .utils import ask_local_llm


class AIWriter:
    """
    根据实验记录的内容撰写文章
    """

    def __init__(
        self,
        logger,
        model: str = "qwen3:14b",
        records: list = [],
        output_dir: str = "/app/temp",
    ):
        self.searcher = PubMedSearcher()

        self.logger = logger
        self.model = model
        self.records = records
        self.output_dir = output_dir
        self.logging_len = 80
        self.records_en = []

        self.papers = []
        self.methods = []
        self.results = []
        self.discussions = []
        self.introductions = []
        self.title = ""
        self.keywords = ""

        self.__init()

    def __init(self):
        if not os.path.exists(self.output_dir):
            os.mkdir(self.output_dir)

        # 数据准备
        self.translate_records()

    def log_handler(self, msg, msg_type: str = "info"):
        """
        统一的日志处理函数
        """
        n = int((self.logging_len - len(msg)) / 2)
        if hasattr(self.logger, msg_type):
            getattr(self.logger, msg_type)(
                f"{'-'*n}{msg}{'-'*(self.logging_len - len(msg) - n)}"
            )
        else:
            self.logger.info(f"{'-'*n}{msg}{'-'*(self.logging_len - len(msg) - n)}")

    def contains_chinese(self, text):
        """检查文本是否包含中文"""
        chinese_pattern = re.compile(r"[\u4e00-\u9fff]")
        return bool(chinese_pattern.search(str(text)))

    def translate_records(self):
        """翻译记录(两层嵌套)"""
        self.log_handler("Translate Records to English")
        if len(self.records) > 0:
            for record in self.records:
                item = {}
                for key, value in record.items():
                    if self.contains_chinese(value):
                        translator = Translator(to_lang="en", from_lang="zh")
                        item[key] = translator.translate(value)
                    else:
                        item[key] = value

                self.records_en.append(item)

    def auto_generate(self):
        """
        自动执行生成论文的工作流
        """
        self.methods_generator()
        self.results_generator()
        self.discussion_generator()
        self.introduction_generator()
        self.abstract_generator()
        self.reference_generator()
        self.title_generator()
        self.keywords_generator()

        print(self.keywords)
        print(self.title)

    def save2file(self, filename, ctx):
        """
        使用写模式,将内容保存到文件
        """
        self.log_handler(f"Save context to '{filename}' file")
        with open(f"{self.output_dir}/{filename}", "w", encoding="utf-8") as f:
            f.write(ctx)

    def paper_searcher(self, query_string, paper_num: int = 20, times: int = 5):
        """文献搜索器

        Args:
            paper_num (int, optional): _description_. Defaults to 10.
            times (int, optional): _description_. Defaults to 5.
        """
        c = 0
        papers = []
        pmids = []

        while c < times and len(papers) < paper_num:
            item_papers = self.searcher.search_papers(query_string, paper_num)
            if len(item_papers) > 0:
                for paper in item_papers:
                    if paper.get("pmid") not in pmids:
                        papers.append(paper)
                        pmids.append(paper.get("pmid"))
            c += 1

        self.papers = papers

    def title_generator(self, max_len=15):
        """标题生成器
        根据提供的内容凝练出标题
        Args:
            ctx (_type_): _description_
        """
        title = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to generate a title based on the content provided by the user.",
                },
                {
                    "role": "system",
                    "content": f"""
                - The return result is within {max_len} words.
                - You don't have to count the number of words in the response.
                - Output only plain text content.
                """,
                },
                {
                    "role": "user",
                    "content": self.abstract,
                },
            ],
            self.model,
        )
        self.title = title

    def keywords_generator(self):
        """
        根据Abstract生成Keywords
        """
        keywords = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to extract keyword information from the content provided by user.",
                },
                {
                    "role": "system",
                    "content": """
                - No more than 5 keywords.
                - This should be a string that will be used to find papers with semantically similar content.
                """,
                },
                {
                    "role": "user",
                    "content": self.abstract,
                },
            ],
            self.model,
        )
        self.keywords = keywords

    def abstract_generator(self):
        """
        根据Introduction生成Abstract
        """
        self.log_handler("Generage Abstract")
        abstract = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to generate a paper abstract based on the research content provided by the user.",
                },
                {
                    "role": "system",
                    "content": """
                - The return result is within 300 words.
                - You don't have to count the number of words in the response.
                - Output only plain text content.
                - The abstract focuses on describing the research results and significance.
                """,
                },
                {
                    "role": "user",
                    "content": "\n\n".join(self.introductions),
                },
            ],
            self.model,
        )
        self.abstract = abstract
        self.save2file("Abstract.md", abstract)

    def introduction_generator(self):
        """
        三段式结构:
        1. 研究背景与意义(文献检索)
        2. 前人研究进展与不足
        3. 本研究的结果与意义
        """

        self.log_handler("Generate Introduction")

        # 先对结果部分进行总结,再提取关键词
        result_summary = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to summarize the research background and significance of users based on the content provided by user.",
                },
                {
                    "role": "system",
                    "content": """
                - The return result is within 200 words.
                - You don't have to count the number of words in the response.                
                """,
                },
                {
                    "role": "user",
                    "content": "\n\n".join(self.results),
                },
            ],
            self.model,
        )

        # 根据总结生成文献检索
        query_string = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to extract keyword information from the content provided by users for literature retrieval.",
                },
                {
                    "role": "system",
                    "content": """
                - No more than 3 keywords.
                - Use OR between different keywords.
                - Respond only with the PubMed search query and nothing else.
                - This should be a string that will be used to find papers with semantically similar content.
                """,
                },
                {
                    "role": "user",
                    "content": result_summary,
                },
                {
                    "role": "user",
                    "content": "Please come up with a search query to find relevant papers on PubMed.",
                },
            ],
            self.model,
        )

        self.paper_searcher(query_string, paper_num=20)

        # 根据检索到的文献生成研究背景与相关工作
        paper_abstracts = [
            f"{paper['abstract']} (PubMed {paper['pmid']})" for paper in self.papers
        ]

        background = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to generate research background based on the context provided by user.",
                },
                {
                    "role": "system",
                    "content": """Please generate content strictly in accordance with the following requirements: 
                - Only one paragraph should be generated, no multiple paragraphs.
                - The text must not exceed 300 words.
                - The text must not contain titles or lists.
                - The reference documents provided by users contain PubMed ID information such as "(PubMed 34662886)", If the text generates a reference to the relevant literature, the PubMed ID should be inserted into the text.
                - You don't have to count the number of words in the response.
                """,
                },
                {
                    "role": "user",
                    "content": f"""The references are listed below: {"\n".join(paper_abstracts)} \n\n My research results are: {result_summary}. Please generate a text describing the background.""",
                },
            ],
            self.model,
        )

        literature_review = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to generate Literature Review in the field based on literature and research provided by users.",
                },
                {
                    "role": "system",
                    "content": """Please generate content strictly in accordance with the following requirements: 
                - Only one paragraph should be generated, no multiple paragraphs.
                - The text must not exceed 300 words.
                - The text must not contain titles or lists.
                - The references provided by the user have a PubMed ID, such as "(PubMed 34662886)". If the text generates a reference to the relevant literature, the PubMed ID should be inserted into the text.
                - You don't have to count the number of words in the response.
                """,
                },
                {
                    "role": "user",
                    "content": f"""The references are listed below: {"\n".join(paper_abstracts)} \n\n My research results are: {result_summary}. Please generate a text describing the relevant work.""",
                },
            ],
            self.model,
        )

        self.introductions = [background, literature_review, result_summary]
        self.save2file("Introduction.md", "\n\n".join(self.introductions))

    def discussion_generator(self):
        """生成结果(结论)部分
        平铺式结构,对每个实验结果进行讨论
        TODO: 查阅参考文献
        """
        self.log_handler("Generate Discussion")
        for result in self.results:
            ctx = ask_local_llm(
                [
                    {
                        "role": "system",
                        "content": "You are a Academic Writing Specialist. Your goal is to take the content provided by users and write it into the Discussion section of an academic paper.\n",
                    },
                    {
                        "role": "system",
                        "content": """The requirements for the answer content are as follows: 
                        1.  Only one paragraph should be generated, no multiple paragraphs.  
                        2.  The text must not exceed 300 words.  
                        3.  Do not include any content unrelated to what the user has provided.
                        4.  You don't have to count the number of words in the response.
                        """,
                    },
                    {"role": "user", "content": result},
                    {
                        "role": "user",
                        "content": "Now please generate a text for me and put forward some guiding suggestions for future work, strictly in line with the requirements of the Discussion section of an academic paper.",
                    },
                ],
                self.model,
            )
            self.discussions.append(ctx)

        self.save2file(f"Discussion.md", "\n\n".join(self.discussions))

    def results_generator(self):
        """生成结果(结论)部分
        平铺式结构,对每个实验结果进行描述
        TODO: 解析每个结果的附件,并从中提炼关键信息
        """
        self.log_handler("Generate Results")
        for index, record in enumerate(self.records_en):
            ctx = ask_local_llm(
                [
                    {
                        "role": "system",
                        "content": "You are a Academic Writing Specialist. Your goal is to take the content provided by users and write it into the Results section of an academic paper.\n",
                    },
                    {
                        "role": "system",
                        "content": "The requirements for the answer content are as follows: 1.  Only one paragraph should be generated, no multiple paragraphs.  2.  The text must not exceed 300 words.  3.  Do not include any content unrelated to what the user has provided.\n",
                    },
                    {
                        "role": "system",
                        "content": """
                        - You can optimize the text provided by the user, but do not modify the data in the content.
                        - Do not hallucinate results that don't exist.
                        - Make sure you clearly and numerically report experimental results in the results section.
                        """,
                    },
                    {
                        "role": "user",
                        "content": f"""Result: {record["result"]} \n Conclusion: {record["conclusion"]}""",
                    },
                    {
                        "role": "user",
                        "content": "Now please generate a text for me, strictly in line with the requirements of the Results section of an academic paper.",
                    },
                ],
                self.model,
            )

            self.results.append(ctx)

        # 整体润色与优化(合并意思相近的内容)
        results_ctx = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to take the content provided by users and write it into the Results section of an academic paper.\n",
                },
                {
                    "role": "system",
                    "content": f"The requirements for the answer content are as follows: 1. The text must not exceed {300*len(self.methods)} words.  2. Do not include any content unrelated to what the user has provided.\n",
                },
                {
                    "role": "system",
                    "content": """
                    - This chapter is mainly a description of the results, and does not involve any experimental method related content.
                    - Make sure you clearly report precise mathematical equations in the Results section and the precise methodology.
                    - The text must not contain titles or lists.
                    - You don't have to count the number of words in the response.
                    - Use "\n\n" segmentation between paragraphs.
                    """,
                },
                {"role": "user", "content": "\n\n".join(self.results)},
                {
                    "role": "user",
                    "content": f"I provided {len(self.results)} descriptions of results. Please check for any duplication between them, if any, merge the duplicate statements and polish the overall description. Strictly in line with the requirements of the Results section of an academic paper.",
                },
            ],
            self.model,
        )
        self.save2file(f"Results.md", results_ctx)

    def methods_generator(self):
        """生成方法部分
        平铺式结构,对每个实验方法进行描述
        分条目生成方法部分
        """
        self.log_handler("Generate Methods")
        for index, record in enumerate(self.records_en):
            ctx = ask_local_llm(
                [
                    {
                        "role": "system",
                        "content": "You are a Academic Writing Specialist. Your goal is to take the content provided by users and write it into the Methods section of an academic paper.\n",
                    },
                    {
                        "role": "system",
                        "content": """
                        The requirements for the answer content are as follows: 
                        1.  Only one paragraph should be generated, no multiple paragraphs.  
                        2.  The text must not exceed 200 words.  
                        3.  The text must not contain titles or lists.
                        4.  Do not include any content unrelated to what the user has provided.
                        """,
                    },
                    {
                        "role": "system",
                        "content": """
                        - What we do. Why we do it. All described using the general Formalism introduced in the Problem Setting and building on top of the concepts / foundations introduced in related work.
                        - If mathematical formulas are involved, use standard LaTeX formatting.
                        - You don't have to count the number of words in the response.
                        """,
                    },
                    {"role": "user", "content": record["step"]},
                    {
                        "role": "user",
                        "content": "Now please generate a text for me, strictly in line with the requirements of the Methods section of an academic paper.",
                    },
                ],
                self.model,
            )

            self.methods.append(ctx)

        # 整体润色与优化(合并意思相近的内容)
        methods_ctx = ask_local_llm(
            [
                {
                    "role": "system",
                    "content": "You are a Academic Writing Specialist. Your goal is to take the content provided by users and write it into the Methods section of an academic paper.\n",
                },
                {
                    "role": "system",
                    "content": f"The requirements for the answer content are as follows: 1. The text must not exceed {200*len(self.methods)} words.  2. Do not include any content unrelated to what the user has provided.\n",
                },
                {
                    "role": "system",
                    "content": """
                    - This chapter mainly introduces the experimental method, which should generally describe the parameters and components of the method. .
                    - Make sure you clearly report precise mathematical equations in the methods section and the precise methodology.
                    - You don't have to count the number of words in the response.
                    - Use "\n\n" segmentation between paragraphs.
                    """,
                },
                {"role": "user", "content": "\n\n".join(self.methods)},
                {
                    "role": "user",
                    "content": f"I provided {len(self.methods)} descriptions of the method. Please check for any duplication between them, if any, merge the duplicate statements and polish the overall description. Strictly in line with the requirements of the Methods section of an academic paper.",
                },
            ],
            self.model,
        )

        self.save2file(f"Methods.md", methods_ctx)

    def reference_generator(self):
        """
        TODO：遍历文章内容, 生成参考文献
        """
        bibliography = "\\begin{thebibliography}{99}\n"
        for paper in self.papers:
            pmid = paper.get("pmid", "")
            authors = ", ".join(paper.get("authors", []))
            title = paper.get("title", "")
            journal = paper.get("journal", "")
            year = paper.get("year", "")

            bibliography += f"""\\bibitem{{PubMed{pmid}}} {authors}. {title}. {journal}. {year}. (PubMed ID: {pmid}) \n"""

        bibliography += "\\end{thebibliography}"

        self.save2file(f"Bibliography.tex", bibliography)
